package org.stayfool.auth.shiro;

import org.apache.shiro.authc.AuthenticationInfo;
import org.apache.shiro.authc.AuthenticationToken;
import org.apache.shiro.authc.ExcessiveAttemptsException;
import org.apache.shiro.authc.credential.HashedCredentialsMatcher;
import org.apache.shiro.cache.Cache;
import org.apache.shiro.cache.CacheManager;
import org.springframework.util.Assert;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * Created by jych1 on 2016/1/1.
 */
public class RetryLimitHashedCredentialsMatcher extends HashedCredentialsMatcher {

	private Cache<String, AtomicInteger> passwordRetryCache;
	private int maxRetryCount = 5;

	public void setCacheManager(CacheManager cacheManager) {
		passwordRetryCache = cacheManager.getCache("pwdRetryCache");
	}

	public void setMaxRetryCount(int maxRetryCount) {
		this.maxRetryCount = maxRetryCount;
	}

	@Override
	public boolean doCredentialsMatch(AuthenticationToken token, AuthenticationInfo info) {

		if (passwordRetryCache == null)
			return super.doCredentialsMatch(token, info);

		Assert.notNull(passwordRetryCache, "cachemanager should not be null");

		String username = (String) token.getPrincipal();
		// retry count + 1
		AtomicInteger currentRetryCount = passwordRetryCache.get(username);
		if (currentRetryCount == null) {
			currentRetryCount = new AtomicInteger(0);
			passwordRetryCache.put(username, currentRetryCount);
		}
		if (currentRetryCount.incrementAndGet() > maxRetryCount) {
			// if retry count > 5 throw
			throw new ExcessiveAttemptsException();
		}

		boolean matches = super.doCredentialsMatch(token, info);

		if (matches) {
			// clear retry count
			passwordRetryCache.remove(username);
		}
		return matches;
	}
}
